import torch

def softmax(a:torch.tensor):
    rowmax = torch.max(a, dim=-1, keepdim=True)[0]
    p = torch.exp(a - rowmax)
    l = torch.sum(p, dim=-1, keepdim=True)
    return p / l

m = torch.rand(3, 5)

rowmax = torch.max(m, dim=-1, keepdim=True)

print(m)
print(rowmax[0])

exp = torch.exp(m - rowmax[0])
print(exp)

